llama : add high-throughput mode #14363
Conversation
github-actions
bot
added
examples
ggml
|
Right now I am comparatively less busy with my PhD so it would be a good time for me to write CUDA code that is still missing, if there is any. |
|
For now, these are the necessary CUDA changes:
// old
// q: [n_embd_k, n_batch, n_head, 1]
// k: [n_embd_k, n_kv, n_head_kv, 1]
// v: [n_embd_v, n_kv, n_head_kv, 1] !! not transposed !!
// mask: [n_kv, n_batch_pad, 1, 1] !! n_batch_pad = GGML_PAD(n_batch, GGML_KQ_MASK_PAD) !!
// res: [n_embd_v, n_head, n_batch, 1] !! permuted !!
GGML_API struct ggml_tensor * ggml_flash_attn_ext(
...);
// new - supports `n_seq` dimension:
// q: [n_embd_k, n_batch, n_head, n_seq]
// k: [n_embd_k, n_kv, n_head_kv, n_seq]
// v: [n_embd_v, n_kv, n_head_kv, n_seq] !! not transposed !!
// mask: [n_kv, n_batch_pad, n_seq, 1] !! n_batch_pad = GGML_PAD(n_batch, GGML_KQ_MASK_PAD) !!
// res: [n_embd_v, n_head, n_batch, n_seq] !! permuted !!
GGML_API struct ggml_tensor * ggml_flash_attn_ext(
...);CPU might also need to be extended (not sure yet)
Edit: the CPU versions of |
ggerganov
force-pushed
the
gg/llama-high-throughput
branch
from
ab2a2bb to
1b74b9d
Compare
ggerganov
force-pushed
the
gg/kv-cache-use-set-rows
branch
3 times, most recently
from
c246784 to
06bb08a
Compare
ggerganov
force-pushed
the
gg/kv-cache-use-set-rows
branch
3 times, most recently
from
82277da to
4534123
Compare
ggerganov
force-pushed
the
gg/kv-cache-use-set-rows
branch
from
2f577c5 to
30b4d4e
Compare
ggerganov
force-pushed
the
gg/llama-high-throughput
branch
from
6179578 to
dfceb01
Compare
ggerganov
force-pushed
the
gg/llama-high-throughput
branch
2 times, most recently
from
eb5856c to
ee0f729
Compare
ggerganov
force-pushed
the
gg/llama-high-throughput
branch
from
ee0f729 to
deae7cd
Compare
ggerganov
force-pushed
the
gg/llama-high-throughput
branch
2 times, most recently
from
988d0cd to
dbcfcaa
Compare
src/llama-kv-cache-unified.cpp
Outdated
| v_cells[s].resize(kv_size); | ||
| } | ||
|
|
||
| // by default, all sequence ids are mapped to the 0th virtual sequence |
There was a problem hiding this comment.
I'd like to understand the purpose of virtual sequences.
- Is it to make the unified cache not unified?
- Should it be a separate cache type instead?
- why is
n_seq_virta number and not aboolof whether or not the cache is unified?- Is it to eventually allow
n_seq_max % n_seq_virt == 0for a partially-unified cache?
- Is it to eventually allow
Are virtual sequences intended to be used with other types of caches eventually (e.g. recurrent)?- The concept here seems specific to the self-attention KV cache (unless I'm misunderstanding).
There was a problem hiding this comment.
Today I found a better term instead of "virtual sequences": "streams". So I'll use "streams" here and will update the code later today or tomorrow.
Is it to make the unified cache not unified?
Roughly yes. The user will be able to select between unified (i.e. single stream) or non-unified (multiple streams). Each mode has advantages in different scenarios. Single stream is good when the sequences share large common prefixes. Multiple streams are good when the sequences are mostly or completely independent from each other.
The first iteration will support 1 stream (i.e. same as master, vanilla unified KV cache) and n_seq_max streams. The latter means that each sequence id is assigned to a separate stream.
In theory, we could assign multiple sequence ids to the same stream to get a partially-unified KV cache, but this would need extra work and it might not have any useful applications. So out of scope for now.
Should it be a separate cache type instead?
There is too much similar logic. Still thinking about it, but most likely it will end up in the same cache type.
The concept here seems specific to the self-attention KV cache (unless I'm misunderstanding)
Yes.
src/llama-batch.h
Outdated
| // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids | ||
| llama_ubatch split_equal(uint32_t n_ubatch, bool sequential); |
There was a problem hiding this comment.
Why are sequential seq_ids required when virtual sequences are used?
Is it because a contiguous (along the virtual sequence dimension) slice of the KV cache is used?
I wonder if there could be a way to avoid this requirement with ggml_get_rows and/or ggml_mul_mat_id. Might not be worth the extra indirection, though.
There was a problem hiding this comment.
Why are sequential seq_ids required when virtual sequences are used?
Is it because a contiguous (along the virtual sequence dimension) slice of the KV cache is used?
Yes, we make a view of the KV cache across the streams here:
llama.cpp/src/llama-kv-cache-unified.cpp
Lines 976 to 992 in dbcfcaa
The ns var is the number of streams that participate in the current ubatch. Their stream indices range from [s0, s1].
I wonder if there could be a way to avoid this requirement with ggml_get_rows and/or ggml_mul_mat_id. Might not be worth the extra indirection, though.
It should be possible. But I'm not sure if it would be worth - both in performance and in complexity. We can explore though.
src/llama-kv-cache-unified.cpp
Outdated
| @@ -45,7 +46,7 @@ llama_kv_cache_unified::llama_kv_cache_unified( | |||
| auto it = ctx_map.find(buft); | |||
| if (it == ctx_map.end()) { | |||
| ggml_init_params params = { | |||
| /*.mem_size =*/ size_t(2u*n_layer_cache*ggml_tensor_overhead()), | |||
| /*.mem_size =*/ size_t(2u*(1 + n_seq_virt)*n_layer_cache*ggml_tensor_overhead()), | |||
There was a problem hiding this comment.
Is the 1 + intended? Why was it added?
There was a problem hiding this comment.
For the per-stream views of the KV cache:
llama.cpp/src/llama-kv-cache-unified.cpp
Lines 125 to 133 in dbcfcaa
These are used to implement the llama_memory_seq_cp(). This operation is no longer just assigning ids - it performs actual copy of the buffers in memory when we use multiple streams. Using these helper views, the operation is quite simple to implement:
llama.cpp/src/llama-kv-cache-unified.cpp
Lines 289 to 329 in dbcfcaa
Though we cannot copy partial sequences when using multiple streams.
src/llama-batch.cpp
Outdated
| // accept only increasing sequence ids | ||
| if (sequential) { | ||
| add = add && (cur_seq_set.empty() || batch.seq_id[i][0] == last_seq_id + 1); | ||
| } |
There was a problem hiding this comment.
What about decreasing sequence ids? Is the requirement that they are increasing, or that the included seq_ids should be in a contiguous range?
(decreasing sequence ids might not really happen often in practice though)
There was a problem hiding this comment.
Decreasing would also work - we just need continuous range. We can either add this, if there is an elegant way to search for this. Or we add some batch pre-processing step to move the complexity at a higher level. Or just delegate it to the user by warning when the batch is not arranged optimally.
ggerganov
force-pushed
the
gg/llama-high-throughput
branch
from
dbcfcaa to
33dcc3c
Compare
|
I did some more performance testing with #14668 but I think it will be necessary to use longer prompts in order to get sensitivity to KV cache changes. |
I think these change slightly improves the user experience - they just have to think about "unified" vs "split (i.e. non-unified)" KV cache and don't have to know about "streams". |
There was a problem hiding this comment.
I think it is ok to keep this gated as an advanced option behind the LLAMA_SET_ROWS environment variable, until all the backends implement support for it.
Other than requiring ggml_set_rows, are there any significant downsides to enabling split KV? I think that if there aren't, it should be the default. If so, it may be better to make it the default already to avoid having to change the command line option later.
I don't think there are. We don't actually have use cases where multiple sequences share large prefixes (neither
Alright. So I will set "kv split" to true, but require the |
|
I did some more performance testing using the new With this PR I get: So there is a ~3x speedup for large contexts that is consistent with the results for 
The first image is master, the second image is this PR. As before, the throughput with |
|
Note: latest commit changes |
For this test, does it make any difference between 1x GPU and 6x GPUs? My understanding is that this test will be dominated by generation time. The prompt processing time would be very short, so pipeline parallelism would likely make small to no difference. |
|
Here are results on M2 Ultra but with shorter prompts: LLAMA_ARG_FLASH_ATTN=1 LLAMA_SET_ROWS=1 LLAMA_ARG_MODEL=./models/llama-3.2-1b-instruct/ggml-model-q4_k.gguf python3 scripts/server-bench-x.py --path_server ./build/bin/llama-server --prompt_source rng-1024-8192
The plots do look a bit unexpected. Wouldn't expect so much chaotic variations in the |
It shouldn't make a meaningful difference for the performance, I just don't have enough VRAM on a single GPU to run the benchmark with a total context size of 1.4M.
Could be when >1 requests happen to be submitted at almost the same time, would maybe make sense to investigate by also plotting the submission/end times of requests. (And since the benchmarking code is pretty new it's also possible that there are just bugs affecting the results.) |
|
Is kv cache quant not supported yet? I've just updated to master, set LLAMA_SET_ROWS=1 and run llama-server with my usual command line but got the following error (even with very small context and just say "hi" from webui): I removed |
|
@rujialiu You can see summary of the supported backends here: #14661. For example, with CUDA, only Btw, how are you using |
|
Thanks for the info! That message appeared before sending any request: BTW: This is a surprisingly usable setup with 12GB VRAM to allow 4 concurrent coding agent with context length 128k. I'm been already using it with cline and able to get some non-trivial jobs done (though should be used with care) :D |
|
There was indeed a bug - will be fixed with #14733. |
|
@ggerganov It looks like mask is not correctly padded with parallel processing: Edit: With |
|
@CISC The diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp
index d8e2086c8..ab6470bdf 100644
--- a/src/llama-memory-hybrid.cpp
+++ b/src/llama-memory-hybrid.cpp
@@ -31,21 +31,21 @@ llama_memory_hybrid::llama_memory_hybrid(
hparams(model.hparams),
mem_attn(new llama_kv_cache_unified(
model,
filter_attn == nullptr ?
[&](int32_t il) { return !hparams.is_recurrent(il); }
: filter_attn,
type_k,
type_v,
v_trans,
offload,
- 1,
+ false,
kv_size,
n_seq_max,
n_pad,
n_swa,
swa_type
)),
mem_recr(new llama_memory_recurrent(
model,
filter_recr == nullptr ?
[&](int32_t il) { return hparams.is_recurrent(il); } |
|
I tested again after #14712 is merged. I can confirm that my use case is working without concurrent access. The VRAM usage is reduced from 11GB->10GB at startup, I've used it for a while with cline without any problem. Then I used claude code to test concurrent access and, after some time, got assertion failed. Here is some log at the end: The command line is almost the same except the server port number. |
|
@rujialiu Thanks for reporting these - very useful information. The CUDA backend needs to be updated to handle quantized KV cache with the new split KV mode (#14756). Btw, are you interested in writing a tutorial for using Claude Code with |
Yes, I'd love to share my experience! Recently I tried a lot of approaches to agentic coding with |
target #14285
Overview
Improve multi-sequence decoding performance by avoiding the cross-sequence attention compute.
Note
To enable this functionality, there is a temporary requirement
LLAMA_SET_ROWS=1to be set in your environment variable. In the future, this will become the default. See below for more info.If you try to use "split KV" cache and haven't added
LLAMA_SET_ROWS=1you will see the following warning:Description
One significant drawback of the unified KV cache is that it leads to performing a lot of unnecessary computation in the attention when the unified buffer is shared between many large independent sequences. The reason is that we have to view this buffer continuously and therefore we end up computing large potions of "cross-sequence attention" which we then simply discard.
With this change, we add option to split the unified KV cache buffer into multiple buffers - one for each sequence. This decouples the sequences from each other and improves the performance and memory usage of the attention when more than one sequence is used. To achieve that, when the batch reaches the attention, we split it into multiple "streams":
llama.cpp/src/llama-graph.cpp
Lines 1035 to 1044 in c96c48c
Each stream has its own KV cache buffer and thus no longer "sees" the rest of the other streams - it attends only to the tokens that belong to the same stream.
With this approach we now have 2 modes:
The new "split" mode is enabled by default. However it requires the
LLAMA_SET_ROWS=1environment variable to be set. Otherwise, a warning will be printed and the context will fallback to "unified" mode. In the future, after there is enoughggml_set_rows()coverage in the backends (#14661) this will become the default mode.To force the old "unified" mode, use
--kv-unifiedCLI arg.API Changes
bool llama_context_params::kv_unified. Default isfalsellama.cpp/include/llama.h
Lines 336 to 340 in fb8150d
Testing
Use
LLAMA_SET_ROWS=1 llama-[command] ...Qwen 2.5 Coder 3B Q8_0, M2 Ultra
Geamma 3 4B Q8_0, M2 Ultra
Using a more real-world example with
llama-parallel:TODO
ggml_soft_max_ext()support for virtual sequencesllama_memory_seq_cpsupport for virtual sequencessplit_equalsupport sequential idsLLAMA_HTbecome regular compute parametern_ctxmeaning (total vs per-sequence)Require(no longer needed)n_embd_v_gqa(il) == constwhen FA is offNext PRs
(split_equal + padding)and stream split[TAG_NO_CACHE_PAD]ggml_set_rows()is fully adoptedllama-parallelto use different RNG seeds for the different clients